[12/17] DLA rank conversion - #2964

What does this PR do?

Type of change: new feature

Add ordered rank and axis preparation stages to the internal DLA sequence. Part 12/17 depends on part 11; merge the series from the bottom upward. Public registration follows in part 16.

Reshape, Expand, Tile and Slice prepare fixed controls for proven four-dimensional layouts, preserve native runtime layouts and allocate private controls for shared/public owners. Copied dimensions and omitted Slice axes are resolved against the native metadata. Unsupported layouts refuse before mutation.

Transpose, normalization, ArgMax, Flatten and axis preparation preserve native public ranks, runtime defaults and opaque/captured owners. Five-dimensional remapping requires proof of the collapsed singleton. Quantizer axes follow imported schema defaults; one-element vectors retain per-axis semantics and scalar controls retain their attributes.

Usage

from modelopt.onnx.graph_surgery.make_dla_compatible import _transform_make_dla_compatible

model = _transform_make_dla_compatible(model)

This slice extends the internal sequence; the public entry point follows in part 16.

Initializer promotion preserves kernels and bias owned by proven native Conv1D paths, including normalization and unary producers. Other convolution conversion paths keep their existing preparation.

GraphCache retains an initializer value-info snapshot for native-rank recovery while its ordinary shape lookup continues to reflect the actual promoted legacy input declaration. Concat axis preparation uses that native snapshot, including legacy IR3 models and absent, concrete or symbolic value-info annotations. Stage and full-pipeline execution cases preserve both positive and negative axes; the previously failing native-rank cases now pass.

The explicit supported contrib QuantizeLinear/DequantizeLinear operators use default axis 1 when the attribute is omitted. Singleton vector scales remain per-axis and receive the rank shift; scalar scales retain per-tensor behavior. Standard-domain schema defaults and older schemas retain their existing rules. Eight new contrib cases fail before the fix; all 24 scalar/vector and explicit/default quantizer cases pass.

Alongside the existing initializer-broadcast proof, rank-one Tile provenance recognizes a supported Reshape whose current target has four entries, including the padded targets and restoring views produced by earlier stages. The existing operator selector excludes opaque and model-local lookalikes. Eight complete-pipeline cases fail before the fix and now preserve values for native rank-two/rank-four Reshape inputs, ArgMax and ReduceSum, with direct and unary Tile paths. Native rank-one controls remain supported; all 25 focused checks pass.

Initializer promotion keeps broadcast operands at native rank when another operand follows a native symbolic data path. Boundary promotion applies the same rule per consumer, preserving fixed runtime biases and legacy initializer declarations while static consumers retain their existing views. Shared static broadcasts, public bias outputs and quantization controls retain their original layouts. Conflicting native and rank-prepared initializer consumers refuse before mutation. Eighteen initializer cases and four boundary cases failed before the corresponding guards; all 302 rank-conversion tests pass. The same converted models execute with batch sizes one and two.

Symbolic public-output restoration requires a supported producer path with native rank two/three, or a fixed short Reshape target that the registered stage pads. Static rank-one inputs alone do not establish promotion. Runtime Slice, NonZero and Range outputs keep their native layouts, including unary tails; genuinely padded Reshape outputs retain restoration. Six native-output cases fail without the guard, while two Reshape controls remain valid. All 310 rank-conversion tests pass, including multiple executions of each converted runtime-dependent model.

Rank/axis stages and public restoration now share the positive layout classifier. Early Slice controls use prospective input promotion only when that stage is registered; later registered stages inspect actual layouts. Public restoration has its own native-owner context and retains all refusal checks. The proof recognizes actual initializer/Conv ranks, padded targets and normalization adapters. Native rank-one Unsqueeze paths retain their controls, including Tile; padded Reshape paths still convert. Twenty-eight native-adapter cases and a missing-input-stage case fail without the corresponding guards; twenty promotion controls remain valid. The complete suite exercises all 359 rank-conversion tests.

Flatten planning now refuses a remaining direct Gemm consumer before changing any branch. The ordered matrix stage still handles its supported Gemm paths first; unrelated Flatten nodes keep their existing conversion. Both branch-order regressions fail without the guard and assert byte-identical model state on refusal. All seven focused Flatten checks pass.

The registered Flatten stage now receives rank-preparation context and preserves native internal paths instead of widening their outputs. Four CPU cases expose incorrect downstream Shape values before the fix and now pass; standalone deferred-rank validation retains its previous refusal behavior. Both existing late-invalid-layout controls also pass.

Arithmetic consumers sharing an initializer with structural readers now use one private Unsqueeze view, retaining the original control rank and payload. Slice, Reshape and Tile cases fail before the fix; six checked CPU cases cover both legacy IR3 declarations and modern IR9 models. All 370 earlier focused rank cases pass, and 22 restacked broadcast/index/control integration cases pass. A Windows CPU sample handles 512 shared-control branches with 64 MiB retained payloads in 0.004939 seconds and +1.0547 MiB sampled peak RSS, preserving every original initializer byte and object and creating one shared view (2 ms polling). Two additional real pipeline checks confirm that final inference already rebuilds deferred initializer metadata; no production metadata workaround was added.

Testing

Windows CPU: 2049 passed in 360.69 seconds across the complete restacked graph-surgery and compatibility suite with Python 3.11.18, ONNX 1.23.2 and ONNX Runtime 1.22.0. The public DLA API uses direct ONNX saving; existing initializer conversion and unrelated surgeries are preserved. Commits are signed and all commit hooks pass with the equivalent Linux mypy platform setting. GPU execution, DLA compilation and the complete Linux suite were not run locally.

python -m pytest tests/unit/onnx/graph_surgery tests/unit/onnx/test_gqa_graph_surgery.py tests/unit/onnx/quantization/test_dq_transpose_surgery.py --noconftest -o addopts= -p no:cacheprovider -q --tb=short

Before your PR is "Ready for review"

• Backward compatible: yes; this slice extends internal DLA functionality.
• Third-party code/new runtime dependencies: none added; contribution and security guidance reviewed.
• Necessary tests: checked synthetic models and before/after evidence are included.
• Changelog: public feature entry follows in part 16.
• Claude approval: pending for the current commit. This remains a draft.

Additional Information

Base branch: dla/review-11-dla-indexing-reductions-and-attention. Review this slice against that branch.

Merge order: PR1 (merged); Part 2 (withdrawn); then Part 3/17 → Part 4/17 → Part 5/17 → Part 6/17 → Part 7/17 → Part 8/17 → Part 9/17 → Part 10/17 → Part 11/17 → Part 12/17 → Part 13/17 → Part 14/17 → Part 15/17 → Part 16/17 → Part 17/17.

Summary by CodeRabbit

• New Features

  • Expanded DLA model conversion support for a wider range of operations, tensor shapes, and layouts, including slicing, transposition, normalization, indexing, and shape adjustments.
  • Improved handling of model inputs and outputs with symbolic or partially known dimensions.

• Bug Fixes

  • Improved preservation of native output shapes and operation behavior during conversion.
  • Added clearer handling of unsupported or ambiguous model structures.